import pybullet as p
import pybullet_data
import math
from typing import List, Tuple, Dict, Optional
import numpy as np
import alphashape
from scipy.spatial import ConvexHull
import trimesh


p.connect(p.GUI)
p.setAdditionalSearchPath(pybullet_data.getDataPath())

#加载桁轨
obstacle = p.loadURDF("./333333/urdf/3333331.urdf", [1.2885, 2.986, 0], [0, 0, 0, 1],useFixedBase=True)


robot_offset_x = 0
robot_offset_y = 0
robot_offset_z = 0  # 障碍物顶部
robot_position = [robot_offset_x, robot_offset_y, robot_offset_z]

# 机械臂旋转180度（倒装）- 绕Y轴旋转180度
robot_orientation = p.getQuaternionFromEuler([0, math.pi, 0])  # 绕Y轴旋转180度

# 加载机械臂（倒装）
robot_id = p.loadURDF("elfin10_l.urdf", robot_position, robot_orientation)
end_effector = p.loadURDF("./Endeffector/urdf/Endeffector.urdf", [0, 0, 0], [0, 0, 0, 1])


gantry_prismatic_joints = []
num_gantry_joints = p.getNumJoints(obstacle)
for i in range(num_gantry_joints):
    joint_info = p.getJointInfo(obstacle, i)
    if joint_info[2] == p.JOINT_PRISMATIC:
        gantry_prismatic_joints.append(i)
print("桁轨滑动关节索引：", gantry_prismatic_joints)


mount_constraint = p.createConstraint(
    obstacle, gantry_prismatic_joints[1], robot_id, -1,
    p.JOINT_FIXED, [0, 0, 0], [0, 0, 0],
    [0, 0, 0], [0, 0, 0, 1],[0, 0, 0, 1]
)



end_effector_orientation = p.getQuaternionFromEuler([math.pi / 2, math.pi / 8 , 0])
end_effector_constraint = p.createConstraint(robot_id, 6, end_effector, -1,
                                                     p.JOINT_FIXED,
                                                     jointAxis=[0, 0, 0],
                                                     parentFramePosition=[0, 0, 0],
                                                     childFramePosition=[-0.065, 0.25, 0],
                                                     parentFrameOrientation=[0, 0, 0, 1],
                                                     childFrameOrientation=end_effector_orientation)


def get_end_effector_pose() -> Tuple[List[float], List[float]]:
    """获取末端执行器的位置和姿态"""
    # 获取末端链接的状态
    link_state = p.getBasePositionAndOrientation(end_effector)

    position = link_state[0]
    orientation = link_state[1]
    return position, orientation





def load_point_cloud_from_file(point_cloud_file: str, p_num) -> Optional[np.ndarray]:
    """
    从文件读取工件点云数据

    Args:
        point_cloud_file: 点云文件路径

    Returns:
        点云数据numpy数组，形状为(N, 3)
    """
    try:
        print(f"正在读取点云文件: {point_cloud_file}")

        # 读取点云数据
        point_cloud = np.loadtxt(point_cloud_file)
        if point_cloud.shape[0] > p_num:
            idx = np.random.choice(point_cloud.shape[0], p_num, replace=False)
            point_cloud = point_cloud[idx]/1000   #毫米->米
        # 确保数据格式正确
        if point_cloud.shape[1] != 3:
            print(f"错误: 点云数据应该是3列(x,y,z)，实际为{point_cloud.shape[1]}列")
            return None

        print(f"成功读取点云数据: {point_cloud.shape[0]} 个点")
        print(f"点云范围: X[{point_cloud[:, 0].min():.3f}, {point_cloud[:, 0].max():.3f}], "
              f"Y[{point_cloud[:, 1].min():.3f}, {point_cloud[:, 1].max():.3f}], "
              f"Z[{point_cloud[:, 2].min():.3f}, {point_cloud[:, 2].max():.3f}]")

        return point_cloud

    except FileNotFoundError:
        print(f"错误: 找不到点云文件 {point_cloud_file}")
        return None
    except Exception as e:
        print(f"读取点云文件时发生错误: {e}")
        return None


def create_mesh_from_points(points: np.ndarray, position: List[float] = [0, 0, 0],
                            orientation: List[float] = [0, 0, 0, 1]) -> Optional[int]:
    """
    用三角网格重建算法

    Args:
        points: 点云数据
        position: 工件在仿真空间中的位置
        orientation: 工件的姿态（四元数）

    Returns:
        创建的工件物体ID
    """
    try:

        alpha_shape = alphashape.alphashape(points, alpha=10)  # alpha需调参
        # 3. 转成三角网格数据
        mesh = trimesh.Trimesh(vertices=np.array(alpha_shape.vertices),
                               faces=np.array(alpha_shape.faces))



        # 创建 mesh 物体
        collision_shape = p.createCollisionShape(
            shapeType=p.GEOM_MESH,
            vertices=mesh.vertices.astype(np.float32),
            indices=mesh.faces.flatten().astype(np.int32),
            meshScale=[1, 1, 1]
        )
        visual_shape = p.createVisualShape(
            shapeType=p.GEOM_MESH,
            vertices=mesh.vertices.astype(np.float32),
            indices=mesh.faces.flatten().astype(np.int32),
            meshScale=[1, 1, 1],
            rgbaColor=[0.7, 0.4, 0.2, 1])

        workpiece_id = p.createMultiBody(
            baseMass=0,
            baseCollisionShapeIndex=collision_shape,
            baseVisualShapeIndex=visual_shape,
            basePosition=position,
            baseOrientation=orientation
        )

        print(f"成功创建网格工件，物体ID: {workpiece_id}")
        return workpiece_id

    except Exception as e:
        print(f"创建网格时发生错误: {e}")
        return None


def create_sphere_approximation(points: np.ndarray, sphere_radius: float = 0.01,
                                position: List[float] = [0, 0, 0]) -> List[int]:
    """
    用多个小球体近似点云形状

    Args:
        points: 点云数据
        sphere_radius: 每个球体的半径
        position: 工件整体偏移位置

    Returns:
        创建的球体物体ID列表
    """
    sphere_ids = []

    try:
        # 可选：对点云进行降采样以减少球体数量
        # 这里使用简单的间隔采样，实际可以用体素化或其他方法
        step = max(1, len(points) // 1000)  # 最多1000个球体
        sampled_points = points[::step]

        print(f"使用 {len(sampled_points)} 个球体近似点云")

        for i, point in enumerate(sampled_points):
            # 计算球体在世界坐标系中的位置
            sphere_pos = [point[0] + position[0],
                          point[1] + position[1],
                          point[2] + position[2]]

            # 创建球体碰撞形状
            collision_shape = p.createCollisionShape(
                shapeType=p.GEOM_SPHERE,
                radius=sphere_radius
            )

            # 创建视觉形状
            visual_shape = p.createVisualShape(
                shapeType=p.GEOM_SPHERE,
                radius=sphere_radius,
                rgbaColor=[0.6, 0.2, 0.8, 0.6]  # 紫色半透明
            )

            # 创建球体
            sphere_id = p.createMultiBody(
                baseMass=0,
                baseCollisionShapeIndex=collision_shape,
                baseVisualShapeIndex=visual_shape,
                basePosition=sphere_pos
            )

            sphere_ids.append(sphere_id)

        print(f"成功创建 {len(sphere_ids)} 个球体")
        return sphere_ids

    except Exception as e:
        print(f"创建球体近似时发生错误: {e}")
        return []


points = load_point_cloud_from_file("output.txt", 2000)

create_mesh_from_points(points)

while True:
    # ee_position, ee_orientation = get_end_effector_pose()
    #
    # sphere_visual = p.createVisualShape(
    #     p.GEOM_SPHERE, radius=0.015,
    #     rgbaColor=[0, 0.5, 1, 0.6]
    # )
    # sphere_body = p.createMultiBody(
    #     baseMass=0, baseVisualShapeIndex=sphere_visual,
    #     basePosition=ee_position
    # )


    p.stepSimulation()